{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "df8313a0",
   "metadata": {},
   "source": [
    "# Chapter 9: A Line-by-Line Implementation of Attention and Transformer\n",
    "\n",
    "This chapter covers\n",
    "\n",
    "* Architecture and functionalities of encoders and decoders in Transformers\n",
    "* How the attention mechanism uses query, key, and value to assign weights to elements in a sequence \n",
    "* \tDifferent types of Transformers\n",
    "* Building a Transformer from scratch for language translation\n",
    "\n",
    "Transformers are advanced deep learning models that excel in handling sequence-to-sequence prediction challenges, outperforming older models like recurrent neural networks (RNNs) and convolutional neural networks (CNNs). Their strength lies in effectively understanding the relationships between elements in input and output sequences over long distances, such as two words far apart in the text. Unlike RNNs, Transformers are capable of parallel training, significantly cutting down training times and enabling the handling of vast datasets. This transformative architecture has been pivotal in the development of large language models (LLMs) like ChatGPT, BERT, and T5, marking a significant milestone in AI progress.\n",
    "\n",
    "Prior to the introduction of Transformers in the groundbreaking 2017 paper Attention Is All You Need by a group of Google researchers,  natural language processing (NLP) and similar tasks primarily relied on RNNs, including long short-term memory (LSTM) models. RNNs, however, process information sequentially, limiting their speed due to the inability to train in parallel and struggling with maintaining information about earlier parts of a sequence, thus failing to capture long-term dependencies.\n",
    "\n",
    "The revolutionary aspect of the Transformer architecture is its attention mechanism. This mechanism assesses the relationship between words in a sequence by assigning weights, determining the degree of relatedness in meaning among words based on the training data. This enables models like ChatGPT to comprehend relationships between words, thus understanding human language more effectively. The non-sequential processing of inputs allows for parallel training, reducing training time and facilitating the use of large datasets, thereby powering the rise of knowledgeable LLMs and the current surge in AI advancements.\n",
    "\n",
    "In this chapter, we will delve into building a Transformer from the ground up, based on the paper Attention Is All You Need. The Transformer, once trained, can handle translations between any two languages (such as German to English or English to Chinese). In the next chapter, we’ll focus on training the Transformer developed here to perform English into French translations. \n",
    "To build the Transformer from scratch, we'll explore the inner workings of the self-attention mechanism, including the roles of query, key, and value vectors, and the computation of scaled dot product attention (SDPA). We'll construct an encoder layer by integrating layer normalization and residual connection into a multi-head attention layer and combining it with a feed-forward layer. We’ll then stack six of these encoder layers to form the encoder. Similarly, we'll develop a decoder in the Transformer that is capable of generating translation one token at a time, based on previous tokens in the translation and the encoder's output.\n",
    "\n",
    "This groundwork will equip you to train the Transformer for translations between any two languages. In the next chapter, you’ll learn to train the Transformer using a dataset containing over 47,000 English-to-French translations. You’ll witness the trained model translating common English phrases to French with an accuracy comparable to using Google Translate."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97468550",
   "metadata": {},
   "source": [
    "# 1\tIntroduction to Transformers and Attention\n",
    "## 1.1\tWhat is attention?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7da22ef8",
   "metadata": {},
   "source": [
    "## 1.2\tThe transformer architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "832f3c43",
   "metadata": {},
   "source": [
    "# 2\tBuild an encoder  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd90f9eb",
   "metadata": {},
   "source": [
    "## 2.1. The Attention Mechanism\n",
    "\n",
    "\n",
    "The *attention()* function is defined in the local module as follows:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2ec5222",
   "metadata": {},
   "source": [
    "```python\n",
    "def attention(query, key, value, mask=None, dropout=None):\n",
    "    d_k = query.size(-1)\n",
    "    scores = torch.matmul(query, \n",
    "              key.transpose(-2, -1)) / math.sqrt(d_k)\n",
    "    if mask is not None:\n",
    "        scores = scores.masked_fill(mask == 0, -1e9)\n",
    "    p_attn = nn.functional.softmax(scores, dim=-1)\n",
    "    if dropout is not None:\n",
    "        p_attn = dropout(p_attn)\n",
    "    return torch.matmul(p_attn, value), p_attn\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09d4bdc7",
   "metadata": {},
   "source": [
    "```python\n",
    "from copy import deepcopy\n",
    "class MultiHeadedAttention(nn.Module):\n",
    "    def __init__(self, h, d_model, dropout=0.1):\n",
    "        super().__init__()\n",
    "        assert d_model % h == 0\n",
    "        self.d_k = d_model // h\n",
    "        self.h = h\n",
    "        self.linears = nn.ModuleList([deepcopy(\n",
    "            nn.Linear(d_model, d_model)) for i in range(4)])\n",
    "        self.attn = None\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        if mask is not None:\n",
    "            mask = mask.unsqueeze(1)\n",
    "        nbatches = query.size(0)  \n",
    "        query, key, value = [l(x).view(nbatches, -1, self.h,\n",
    "           self.d_k).transpose(1, 2)\n",
    "         for l, x in zip(self.linears, (query, key, value))]\n",
    "        x, self.attn = attention(\n",
    "            query, key, value, mask=mask, dropout=self.dropout)\n",
    "        x = x.transpose(1, 2).contiguous().view(\n",
    "            nbatches, -1, self.h * self.d_k)\n",
    "        output = self.linears[-1](x)\n",
    "        return output \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca029dbb",
   "metadata": {},
   "source": [
    "```python\n",
    "class PositionwiseFeedForward(nn.Module):\n",
    "    def __init__(self, d_model, d_ff, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.w_1 = nn.Linear(d_model, d_ff)\n",
    "        self.w_2 = nn.Linear(d_ff, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h1 = self.w_1(x)\n",
    "        h2 = self.dropout(h1)\n",
    "        return self.w_2(h2)   \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a568725",
   "metadata": {},
   "source": [
    "## 2.2\tCreate an encoder\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "112ae1cc",
   "metadata": {},
   "source": [
    "```python\n",
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self, size, self_attn, feed_forward, dropout):\n",
    "        super().__init__()\n",
    "        self.self_attn = self_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = nn.ModuleList([deepcopy(\n",
    "        SublayerConnection(size, dropout)) for i in range(2)])\n",
    "        self.size = size  \n",
    "    def forward(self, x, mask):\n",
    "        x = self.sublayer[0](\n",
    "            x, lambda x: self.self_attn(x, x, x, mask))\n",
    "        output = self.sublayer[1](x, self.feed_forward)\n",
    "        return output \n",
    "    \n",
    "class SublayerConnection(nn.Module):\n",
    "    def __init__(self, size, dropout):\n",
    "        super().__init__()\n",
    "        self.norm = LayerNorm(size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "    def forward(self, x, sublayer):\n",
    "        output = x + self.dropout(sublayer(self.norm(x)))\n",
    "        return output  \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c0f149c",
   "metadata": {},
   "source": [
    "```python\n",
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self, features, eps=1e-6):\n",
    "        super().__init__()\n",
    "        self.a_2 = nn.Parameter(torch.ones(features))\n",
    "        self.b_2 = nn.Parameter(torch.zeros(features))\n",
    "        self.eps = eps\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True) \n",
    "        std = x.std(-1, keepdim=True)\n",
    "        x_zscore = (x - mean) / torch.sqrt(std ** 2 + self.eps)\n",
    "        output = self.a_2*x_zscore+self.b_2\n",
    "        return output\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89828266",
   "metadata": {},
   "source": [
    "The encoder consists of N=6 identical encoder layers. The *Encoder* class is defined as follows in the local module: "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e76d313b",
   "metadata": {},
   "source": [
    "```python\n",
    "# Create an encoder\n",
    "from copy import deepcopy\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, layer, N):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList(\n",
    "            [deepcopy(layer) for i in range(N)])\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, mask)\n",
    "            output = self.norm(x)\n",
    "        return output\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a4cec69",
   "metadata": {},
   "source": [
    "# 3\tBuild an encoder-decoder Transformer\n",
    "## 3.1\tCreate a decoder layer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd150da1",
   "metadata": {},
   "source": [
    "```python\n",
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, size, self_attn, src_attn,\n",
    "                 feed_forward, dropout):\n",
    "        super().__init__()\n",
    "        self.size = size\n",
    "        self.self_attn = self_attn\n",
    "        self.src_attn = src_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = nn.ModuleList([deepcopy(\n",
    "        SublayerConnection(size, dropout)) for i in range(3)])\n",
    "\n",
    "    def forward(self, x, memory, src_mask, tgt_mask):\n",
    "        x = self.sublayer[0](x, lambda x: \n",
    "                 self.self_attn(x, x, x, tgt_mask))\n",
    "        x = self.sublayer[1](x, lambda x:\n",
    "                 self.src_attn(x, memory, memory, src_mask))\n",
    "        output = self.sublayer[2](x, self.feed_forward)\n",
    "        return output \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e50fadd",
   "metadata": {},
   "source": [
    "## 3.2\tCreate an encoder-decoder Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85c45299",
   "metadata": {},
   "source": [
    "```python\n",
    "# Create a decoder\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, layer, N):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList(\n",
    "            [deepcopy(layer) for i in range(N)])\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "\n",
    "    def forward(self, x, memory, src_mask, tgt_mask):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, memory, src_mask, tgt_mask)\n",
    "        output = self.norm(x)\n",
    "        return output\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f9204e",
   "metadata": {},
   "source": [
    "```python\n",
    "# An encoder-decoder transformer\n",
    "class Transformer(nn.Module):\n",
    "    def __init__(self, encoder, decoder,\n",
    "                 src_embed, tgt_embed, generator):\n",
    "        super().__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_embed = src_embed\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.generator = generator\n",
    "\n",
    "    def encode(self, src, src_mask):\n",
    "        return self.encoder(self.src_embed(src), src_mask)\n",
    "\n",
    "    def decode(self, memory, src_mask, tgt, tgt_mask):\n",
    "        return self.decoder(self.tgt_embed(tgt), \n",
    "                            memory, src_mask, tgt_mask)\n",
    "\n",
    "    def forward(self, src, tgt, src_mask, tgt_mask):\n",
    "        memory = self.encode(src, src_mask)\n",
    "        output = self.decode(memory, src_mask, tgt, tgt_mask)\n",
    "        return output\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11199cf5",
   "metadata": {},
   "source": [
    "# 4. Put All Pieces Together\n",
    "## 4.1\tDefine a generator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "090eafd8",
   "metadata": {},
   "source": [
    "```python\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, d_model, vocab):\n",
    "        super().__init__()\n",
    "        self.proj = nn.Linear(d_model, vocab)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.proj(x)\n",
    "        probs = nn.functional.log_softmax(out, dim=-1)\n",
    "        return probs  \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0355ede",
   "metadata": {},
   "source": [
    "## 4.2\tCreate a model to translate between two languages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "270bfb7f",
   "metadata": {},
   "source": [
    "```python\n",
    "# create the model\n",
    "def create_model(src_vocab, tgt_vocab, N, d_model,\n",
    "                 d_ff, h, dropout=0.1):\n",
    "    attn=MultiHeadedAttention(h, d_model).to(DEVICE)\n",
    "    ff=PositionwiseFeedForward(d_model, d_ff, dropout).to(DEVICE)\n",
    "    pos=PositionalEncoding(d_model, dropout).to(DEVICE)\n",
    "    model = Transformer(\n",
    "        Encoder(EncoderLayer(d_model,deepcopy(attn),deepcopy(ff),\n",
    "                             dropout).to(DEVICE),N).to(DEVICE),\n",
    "        Decoder(DecoderLayer(d_model,deepcopy(attn),\n",
    "             deepcopy(attn),deepcopy(ff), dropout).to(DEVICE),\n",
    "                N).to(DEVICE),\n",
    "        nn.Sequential(Embeddings(d_model, src_vocab).to(DEVICE),\n",
    "                      deepcopy(pos)),\n",
    "        nn.Sequential(Embeddings(d_model, tgt_vocab).to(DEVICE),\n",
    "                      deepcopy(pos)),\n",
    "        Generator(d_model, tgt_vocab)).to(DEVICE)\n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "    return model.to(DEVICE)\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
